import json

import torch
from transformers.data.processors import InputExample, InputFeatures
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer

train_data_file = "./dataset/KUAKE-QQR_train.json"

input_example_list = []
with open(train_data_file, 'r', encoding='UTF-8') as f:
    json_content = json.load(f)
    for block in json_content:
        input_example = InputExample(guid=block['id'],
                                     text_a=block['query1'],
                                     text_b=block['query2'],
                                     label=block['label'])
        input_example_list.append(input_example)


class QQRDataset(Dataset):
    def __init__(self, input_example_list: list[InputExample]):
        self.input_example_list = input_example_list

    def __getitem__(self, idx):
        input_example = self.input_example_list[idx]
        return input_example.text_a, input_example.text_b, input_example.label

    def __len__(self):
        return len(input_example_list)

class QQRDataset2(Dataset):
    def __init__(self, input_example_list: list[InputExample]):
        self.input_example_list = input_example_list

    def __getitem__(self, idx):
        input_example = self.input_example_list[idx]
        item = dict()
        item['q1'] = input_example.text_a
        item['q2'] = input_example.text_b
        item['label'] = input_example.label
        return item

    def __len__(self):
        return len(input_example_list)

train_dataset = QQRDataset(input_example_list)
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
for item in train_dataloader:
    print(item)
    print(item[0])
    print(item[1])
    print(item[2])
    print(type(item))
    break
print("-----")
# 对应与__getitem__返回的三个参数
for q1, q2, label in train_dataloader:
    print(q1)
    print(q2)
    print(label)
    break
print("-----")
train_dataset2 = QQRDataset2(input_example_list)
train_dataloader2 = DataLoader(train_dataset2, batch_size=32, shuffle=True)
for item in train_dataloader2:
    print(item)
    print(item['q1'])
    print(item['q2'])
    print(item['label'])
    print(type(item))
    break

